import pandas as pd
import numpy as np
import yfinance as yf
import os
from datetime import datetime
import matplotlib.pyplot as plt
import pickle
from pandas.plotting import register_matplotlib_converters
register_matplotlib_converters()


# 创建数据存储目录
DATA_DIR = "stock_data"
os.makedirs(DATA_DIR, exist_ok=True)

# 股票列表（添加交易所后缀）
STOCKS = {
    '600588.SS': '用友网络',
    '600120.SS': '浙江东方',
    '000810.SZ': '创维数字',
    '603893.SS': '瑞芯微',
    '600584.SS': '长电科技',
    '300251.SZ': '光线传媒'
}

def save_var(obj, filename):
    with open(filename, 'wb') as f:  # 二进制写入模式
        pickle.dump(obj, f)

def load_var(filename):
    with open(filename, 'rb') as f:  # 二进制读取模式
        return pickle.load(f)

def get_stock_data(ticker):
    """获取股票数据（优先从本地加载，不存在则从yfinance下载）"""
    file_path = os.path.join(DATA_DIR, f"{ticker}.csv")
    print(file_path)
    # 尝试从本地加载
    if os.path.exists(file_path):
        try:
            return load_var(file_path) 
        except Exception as e:
            print(f"加载本地数据失败，重新下载: {ticker}, 错误: {str(e)}")
    
    # 从yfinance下载
    print(f"下载数据: {ticker}")
    try:
        # 增加数据量确保有足够数据计算指标
        df = yf.download(ticker, period="4y", interval="1mo", auto_adjust=True)
        if df.empty:
            print(df)
            print(f"无法下载 {ticker} 的数据")
            return None
            
        # 保存到本地
        save_var(df, file_path)
        #df.to_csv(file_path)
        print(f"已保存 {ticker} 数据到本地")
        return df.sort_index(ascending=True)
    except Exception as e:
        print(f"下载 {ticker} 失败: {str(e)}")
        return None


def calculate_technical_indicators(df):
    """计算技术指标"""
    if df is None or df.empty:
        return None
    
    # 确保有足够的数据
    if len(df) < 35:  # 需要至少35个月数据计算30月线
        print(f"数据不足，无法计算技术指标 (只有 {len(df)} 个月数据)")
        return None
    
    # 计算典型价格
    tp = (df['High'] + df['Low'] + df['Close']) / 3
    
    # 计算CCI (20周期) - 修复计算方式
    cci_period = 20
    sma_tp = tp.rolling(cci_period).mean()
    # 使用简单函数计算MAD
    mad = tp.rolling(cci_period).apply(lambda x: np.mean(np.abs(x - np.mean(x))), raw=False)
    df['CCI'] = (tp - sma_tp) / (0.015 * mad)
    
    # 计算MACD
    exp12 = df['Close'].ewm(span=12, adjust=False).mean()
    exp26 = df['Close'].ewm(span=26, adjust=False).mean()
    df['MACD'] = exp12 - exp26
    df['Signal'] = df['MACD'].ewm(span=9, adjust=False).mean()
    
    # 计算移动平均线
    df['MA5'] = df['Close'].rolling(window=5).mean()
    df['MA10'] = df['Close'].rolling(window=10).mean()
    df['MA20'] = df['Close'].rolling(window=20).mean()
    df['MA30'] = df['Close'].rolling(window=30).mean()
    
    # 只保留有足够数据计算所有指标的部分
    return df.dropna()

def check_conditions(df):
    """检查技术指标条件并返回数值"""
    if df is None or len(df) < 1:
        return {
            'CCI': False, 
            'MACD': False, 
            'GoldenCross': False,
            'Error': '数据不足',
            'MA5': None,
            'MA10': None,
            'MA20': None,
            'MA30': None,
            'MACD_value': None,
            'CCI_value': None
        }
    
    try:
        # 获取最后两个月的值
        last_row = df.iloc[-1]
        prev_row = df.iloc[-2] if len(df) >= 2 else None
        
        # 获取指标值（无论条件是否满足）
        ma5_value = last_row.get('MA5', None).item()
        ma10_value = last_row.get('MA10', None).item()
        ma20_value = last_row.get('MA20', None).item()
        ma30_value = last_row.get('MA30', None).item()
        macd_value = last_row.get('MACD', None).item()
        cci_value = last_row.get('CCI', None).item()

        # 检查CCI条件
        cci_condition = cci_value > 100 if cci_value is not None else False
        #print(f"CCI: {cci_value}, Condition: {cci_condition}")
        # 检查MACD条件
        macd_condition = macd_value > 0 if macd_value is not None else False
        #print(f"MACD: {macd_value}, Condition: {macd_condition}")

        #检查下降是否回头
        #30月线和20月线都在10月线和5月线上面，且5月线和10月线相差不超过5月线的4%
        lhf_panduan = (
            (ma30_value > ma10_value) and
            (ma20_value > ma10_value) and
            (ma5_prev >= ma10_prev) and  
            (ma30_value > ma5_value) and
            (ma20_value > ma5_value)
        )
        print(f"lhf_panduan---------------: {lhf_panduan}")
        # 检查均线金叉条件
        golden_cross = False
        if prev_row is not None and all(x is not None for x in [
            ma5_value, ma10_value, ma20_value, 
            prev_row.get('MA5'), prev_row.get('MA10'), prev_row.get('MA20')
        ]):
            ma5_prev = prev_row['MA5'].item()
            ma10_prev = prev_row['MA10'].item()
            ma20_prev = prev_row['MA20'].item()
            #print(f"Prev: {ma5_prev}, {ma10_prev}, {ma20_prev}")
            
            golden_cross = (
                (ma5_value > ma10_value) and
                (ma10_value > ma20_value) and
                (ma5_prev <= ma10_prev) and  # 5上穿10
                (ma10_prev <= ma20_prev)     # 10上穿20
            )

        # print(f"ma5_value: {ma5_value}")
        # print(f"ma10_value: {ma10_value}")
        # print(f"ma20_value: {ma20_value}")
        # print(f"ma30_value: {ma30_value}")
        # print(f"macd_value: {macd_value}")
        # print(f"cci_value: {cci_value}")
        # print(f"cci_condition: {cci_condition}")
        # print(f"macd_condition: {macd_condition}")
        # print(f"golden_cross: {golden_cross}")


        return {
            'CCI': cci_condition,
            'MACD': macd_condition,
            'GoldenCross': golden_cross,
            'Error': None,
            'MA5': ma5_value,
            'MA10': ma10_value,
            'MA20': ma20_value,
            'MA30': ma30_value,
            'MACD_value': macd_value,
            'CCI_value': cci_value
        }
        
    except Exception as e:
        print(f"检查条件时出错: {str(e)}")
        
        return {
            'CCI': False, 
            'MACD': False, 
            'GoldenCross': False,
            'Error': str(e),
            'MA5': None,
            'MA10': None,
            'MA20': None,
            'MA30': None,
            'MACD_value': None,
            'CCI_value': None
        }


# 安全格式化函数
def safe_format(value, fmt=".2f"):
    """安全地格式化值，处理None情况"""
    if value is None:
        return "N/A"
    try:
        return format(value, fmt)
    except:
        return str(value)

def analyze_stocks():
    """分析所有股票并返回结果"""
    results = []
    
    for ticker, name in STOCKS.items():
        print(f"\n开始分析 {name}({ticker})")
        df = get_stock_data(ticker)
        if df is None:
            print(f"无法获取 {name}({ticker}) 的数据")
            results.append({
                'Ticker': ticker,
                'Name': name,
                'CCI': False,
                'MACD': False,
                'GoldenCross': False,
                'Error': '无数据',
                'MA5': None,
                'MA10': None,
                'MA20': None,
                'MA30': None,
                'MACD_value': None,
                'CCI_value': None
            })
            continue
            
        df_indicators = calculate_technical_indicators(df)
        if df_indicators is None or df_indicators.empty:
            print(f"无法计算 {name}({ticker}) 的技术指标")
            results.append({
                'Ticker': ticker,
                'Name': name,
                'CCI': False,
                'MACD': False,
                'GoldenCross': False,
                'Error': '技术指标计算失败',
                'MA5': None,
                'MA10': None,
                'MA20': None,
                'MA30': None,
                'MACD_value': None,
                'CCI_value': None
            })
            continue
            
        conditions = check_conditions(df_indicators)
        conditions['Ticker'] = ticker
        conditions['Name'] = name
        results.append(conditions)
        
        # 打印当前股票结果
        print(f"{name}({ticker}) 分析结果:")
        print(f"CCI>100: {'是' if conditions['CCI'] else '否'} (当前值: {safe_format(conditions['CCI_value'])}")
        print(f"MACD>0: {'是' if conditions['MACD'] else '否'} (当前值: {safe_format(conditions['MACD_value'], '.4f')})")
        print(f"均线金叉: {'是' if conditions['GoldenCross'] else '否'}")
        print(f"5月线: {safe_format(conditions['MA5'])}, 10月线: {safe_format(conditions['MA10'])}, 20月线: {safe_format(conditions['MA20'])}, 30月线: {safe_format(conditions['MA30'])}")
        
        if conditions['Error']:
            print(f"错误信息: {conditions['Error']}")
    
    return pd.DataFrame(results)

def plot_technical_chart(ticker, name):
    """绘制技术指标图表"""
    print(f"为 {name}({ticker}) 生成图表...")
    df = get_stock_data(ticker)
    if df is None:
        return
        
    df = calculate_technical_indicators(df)
    if df is None or len(df) < 20:
        print(f"无法为 {name}({ticker}) 生成图表，数据不足")
        return
        
    plt.figure(figsize=(14, 10))
    
    # 价格和均线
    plt.subplot(3, 1, 1)
    plt.plot(df.index, df['Close'], label='收盘价', color='black')
    plt.plot(df.index, df['MA5'], label='5月线', color='blue')
    plt.plot(df.index, df['MA10'], label='10月线', color='orange')
    plt.plot(df.index, df['MA20'], label='20月线', color='green')
    plt.plot(df.index, df['MA30'], label='30月线', color='red')
    plt.title(f'{name}({ticker}) 价格与均线')
    plt.legend()
    plt.grid(True)
    # MACD
    plt.subplot(3, 1, 2)
    plt.bar(df.index, df['MACD'], label='MACD', color=np.where(df['MACD'] > 0, 'g', 'r'))
    plt.plot(df.index, df['Signal'], label='信号线', color='orange')
    plt.axhline(0, color='gray', linestyle='--')
    plt.title('MACD指标')
    plt.legend()
    plt.grid(True)
    
    # CCI
    plt.subplot(3, 1, 3)
    plt.plot(df.index, df['CCI'], label='CCI', color='purple')
    plt.axhline(100, color='red', linestyle='--', label='100水平线')
    plt.axhline(0, color='gray', linestyle='--')
    plt.axhline(-100, color='gray', linestyle='--')
    plt.title('CCI指标')
    plt.legend()
    plt.grid(True)
    
    plt.tight_layout()
    plt.savefig(f"{DATA_DIR}/{ticker}_analysis.png", dpi=100)
    plt.close()
    print(f"已保存 {name}({ticker}) 图表")

def main():
    print("开始股票分析...")
    results_df = analyze_stocks()
    
    if results_df.empty:
        print("\n没有获取到任何股票数据")
        return
    
    # 保存结果到CSV
    results_df.to_csv("stock_analysis_results.csv", index=False)
    print("\n分析结果已保存到 stock_analysis_results.csv")
    

    # df = get_stock_data('600584.SS')
    # df = calculate_technical_indicators(df)
    # print( df['MA10'])


    # # 列出所有股票的指标数值
    # print("\n股票技术指标数值:")
    # for idx, row in results_df.iterrows():
    #     print(f"\n{row['Name']}({row['Ticker']}):")
    #     print(f"  5月线: {safe_format(row['MA5'])}")
    #     print(f" 10月线: {safe_format(row['MA10'])}")
    #     print(f" 20月线: {safe_format(row['MA20'])}")
    #     print(f" 30月线: {safe_format(row['MA30'])}")
    #     print(f" MACD值: {safe_format(row['MACD_value'], '.4f')}")
    #     print(f" CCI值: {safe_format(row['CCI_value'])}")
    
    # 为每只股票生成技术图表
    #print("\n生成技术分析图表...")
    #for ticker, name in STOCKS.items():
    #    plot_technical_chart(ticker, name)
    #print("图表生成完成，保存在 stock_data 目录")

if __name__ == "__main__":
    main()
